package scanner

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"time"
	"kscan/pkg/poc"

	"github.com/fatih/color"
)

// ScanResult 扫描结果
type ScanResult struct {
	Target      string            `json:"target"`
	PocID       string            `json:"poc_id"`
	PocName     string            `json:"poc_name"`
	Severity    string            `json:"severity"`
	Description string            `json:"description"`
	Matched     bool              `json:"matched"`
	StatusCode  int               `json:"status_code"`
	ResponseTime int64            `json:"response_time"`
	Timestamp   time.Time         `json:"timestamp"`
	Details     map[string]string `json:"details,omitempty"`
}

// ScanReport 扫描报告
type ScanReport struct {
	StartTime    time.Time      `json:"start_time"`
	EndTime      time.Time      `json:"end_time"`
	TotalTargets int            `json:"total_targets"`
	TotalPOCs    int            `json:"total_pocs"`
	TotalScans   int            `json:"total_scans"`
	Vulnerabilities int         `json:"vulnerabilities"`
	Results      []*ScanResult  `json:"results"`
}

// NewScanResult 创建扫描结果
func NewScanResult(target string, p *poc.POC, matched bool, response *poc.Response) *ScanResult {
	return &ScanResult{
		Target:       target,
		PocID:        p.ID,
		PocName:      p.Name,
		Severity:     p.Severity,
		Description:  p.Description,
		Matched:      matched,
		StatusCode:   response.StatusCode,
		ResponseTime: response.Time,
		Timestamp:    time.Now(),
	}
}

// PrintResults 打印扫描结果
func PrintResults(results []*ScanResult, verbose bool) {
	red := color.New(color.FgRed, color.Bold)
	yellow := color.New(color.FgYellow, color.Bold)
	green := color.New(color.FgGreen)
	cyan := color.New(color.FgCyan)
	
	vulnerabilityCount := 0
	
	fmt.Println("\n" + strings.Repeat("=", 80))
	fmt.Println("扫描结果")
	fmt.Println(strings.Repeat("=", 80))
	
	for _, result := range results {
		if result.Matched {
			vulnerabilityCount++
			
			// 根据严重程度设置颜色
			var severityColor *color.Color
			switch result.Severity {
			case "critical":
				severityColor = red
			case "high":
				severityColor = red
			case "medium":
				severityColor = yellow
			case "low":
				severityColor = yellow
			default:
				severityColor = green
			}
			
			fmt.Println()
			severityColor.Printf("[%s] %s\n", result.Severity, result.PocName)
			cyan.Printf("目标: %s\n", result.Target)
			fmt.Printf("POC ID: %s\n", result.PocID)
			fmt.Printf("描述: %s\n", result.Description)
			fmt.Printf("状态码: %d\n", result.StatusCode)
			fmt.Printf("响应时间: %dms\n", result.ResponseTime)
			fmt.Printf("发现时间: %s\n", result.Timestamp.Format("2006-01-02 15:04:05"))
			
			if len(result.Details) > 0 {
				fmt.Println("详细信息:")
				for key, value := range result.Details {
					fmt.Printf("  %s: %s\n", key, value)
				}
			}
		} else if verbose {
			green.Printf("[√] %s - %s (未发现漏洞)\n", result.Target, result.PocName)
		}
	}
	
	fmt.Println("\n" + strings.Repeat("=", 80))
	if vulnerabilityCount > 0 {
		red.Printf("发现 %d 个漏洞!\n", vulnerabilityCount)
	} else {
		green.Println("未发现漏洞")
	}
	fmt.Println(strings.Repeat("=", 80) + "\n")
}

// SaveResults 保存结果到文件
func SaveResults(results []*ScanResult, filename string) error {
	report := &ScanReport{
		EndTime:         time.Now(),
		Results:         results,
		TotalScans:      len(results),
	}
	
	// 统计漏洞数量
	for _, result := range results {
		if result.Matched {
			report.Vulnerabilities++
		}
	}
	
	// 转换为JSON
	data, err := json.MarshalIndent(report, "", "  ")
	if err != nil {
		return err
	}
	
	// 写入文件
	return os.WriteFile(filename, data, 0644)
}